from sqlalchemy import insert, select

from owl_admin.ext import db
from owl_ai.domain.model.knowledge_entity import KnowledgeBaseEntity, KnowledgeFileEntity
from owl_ai.repository.knowledge_po import KnowledgeBasePo, KnowledgeFilePO
from owl_common.sqlalchemy.model import ColumnEntityList


class KnowledgeRepository:
    default_fields = {"id", "name", "description", "type", "vector_store", "chunk_strategy", "embedding_strategy",
                      "status", "create_time", "update_time", "create_by", "update_by"}
    default_columns = ColumnEntityList(KnowledgeBasePo, default_fields, False)
    """
    知识库仓库
    """

    @classmethod
    def insert(cls, knowledge: KnowledgeBaseEntity):
        """
        插入知识库
        :param knowledge: 知识库实体
        :return: 插入的id
        """
        fields = {"name", "description", "type", "vector_store", "chunk_strategy", "embedding_strategy", "status",
                  "create_time", "update_time", "create_by", "update_by"}
        data = knowledge.model_dump(include=fields, exclude_none=True, exclude_unset=True)
        stmt = insert(KnowledgeBasePo).values(data)
        out = db.session.execute(stmt).inserted_primary_key
        return out[0] if out else 0

    @classmethod
    def list(cls):
        """
        获取知识库列表
        :return: 知识库列表
        """
        stmt = select(*cls.default_columns).order_by(KnowledgeBasePo.create_time.desc())
        rows = db.session.execute(stmt).all()
        return [cls.default_columns.cast(row, KnowledgeBaseEntity) for row in rows]

    @classmethod
    def find_by_id(cls, kid: int):
        criteria = [KnowledgeBasePo.id == kid]
        stmt = select(*cls.default_columns).where(*criteria)
        row = db.session.execute(stmt).one_or_none()
        return cls.default_columns.cast(row, KnowledgeBaseEntity) if row else None


class KnowledgeFileRepository:
    """
    知识库文件仓库
    """

    default_fields = {"id", "kid", "file_name", "disk_name", "file_type", "create_time", "update_time", "create_by",
                      "update_by"}

    default_columns = ColumnEntityList(KnowledgeFilePO, default_fields, False)

    @classmethod
    def insert(cls, knowledge_file: KnowledgeFileEntity):
        """
        插入知识库文件
        :param knowledge_file: 知识库文件实体
        :return: 插入的id
        """
        fields = {"kid", "file_name", "disk_name", "file_type", "create_time", "update_time", "create_by", "update_by"}
        data = knowledge_file.model_dump(include=fields, exclude_none=True, exclude_unset=True)
        stmt = insert(KnowledgeFilePO).values(data)
        out = db.session.execute(stmt).inserted_primary_key
        return out[0] if out else 0

    @classmethod
    def list(cls, kid: int):
        """
        获取知识库文件列表
        :return: 知识库文件列表
        """
        criteria = []

        if kid:
            criteria.append(KnowledgeFilePO.kid == kid)

        stmt = select(*cls.default_columns).where(*criteria).order_by(KnowledgeFilePO.create_time.desc())
        rows = db.session.execute(stmt).all()
        return [cls.default_columns.cast(row, KnowledgeFileEntity) for row in rows]

    @classmethod
    def find_by_id(cls, kid: int):
        criteria = [KnowledgeFilePO.id == kid]
        stmt = select(*cls.default_columns).where(*criteria)
        row = db.session.execute(stmt).one_or_none()
        return cls.default_columns.cast(row, KnowledgeFileEntity) if row else None
